#!/usr/bin/python
# -*- coding: utf8 -*-

import os
import uuid
import shutil
import json

from flask import Flask, request, Response

from utils import (server_info, JUDGER_WORKSPACE, TEST_CASE_DIR,
                  token, spj_config, logger, ProblemIOMode)
from exception import TokenVerificationFailed, CompileError, SPJCompileError, JudgeClientError, FileIOError
from compiler import Compiler
from judge_client import JudgeClient

app = Flask(__name__)
DEBUG = os.environ.get("judger_debug") == "1"
app.debug = DEBUG


class InitSubmissionEnv(object):
  def __init__(self, judger_workspace, submission_id):
    self.work_dir = os.path.join(judger_workspace, submission_id)
  
  def __enter__(self):
    try:
      os.mkdir(self.work_dir)
      os.chmod(self.work_dir, 0o711)
    except Exception as e:
      logger.exception(e)
      raise JudgeClientError('Failed to create runtime dir.')
    return self.work_dir
  
  def __exit__(self, exc_type, exc_val, exc_tb):
    if not DEBUG:
      pass
      try:
        shutil.rmtree(self.work_dir)
      except Exception as e:
        logger.exception(e)
        raise JudgeClientError("Failed to clean runtime dir.")


class JudgeServer:
  @classmethod
  def ping(cls):
    data = server_info()
    data["action"] = "ping"
    return data
  
  @classmethod
  def judge(cls, language_config, src, time_limit, memory_limit, test_case_path,
            io_mode=None, is_spj=False, in_file=None, out_file=None):
    if not io_mode:
      io_mode = ProblemIOMode.standard

    if not test_case_path:
      raise JudgeClientError("Invalid test case path.")
    test_case_path = os.path.join(TEST_CASE_DIR, test_case_path)

    compile_config = language_config['compile']
    run_config = language_config['run']
    submission_id = uuid.uuid4().hex
    
    with InitSubmissionEnv(JUDGER_WORKSPACE, str(submission_id)) as dir:
      work_dir = dir

      if is_spj:
        # get spj.cpp from test case dir
        spj_src = os.path.join(test_case_path, "spj.cpp")
        spj = os.path.join(work_dir, "spj.cpp")
        try:
          os.system('cp %s %s'%(spj_src, spj))
          os.chmod(spj, 0o400)
        except Exception:
          raise SPJCompileError("File spj.cpp does not exists.")
        Compiler().compileSrc(spj, work_dir, spj_config, is_spj)
        spj_exe = os.path.join(work_dir, 'spj')
        try:
          os.chmod(spj_exe, 0o500)
        except Exception:
          pass
        
        
      if io_mode == ProblemIOMode.file:
        if src.find(out_file) != -1:
          src = src.replace(out_file, os.path.join(work_dir, out_file), 1)
        else:
          raise FileIOError("File IO Error")
        if src.find(in_file) != -1:
          src = src.replace(in_file, os.path.join(work_dir, in_file), 1)
        else:
          raise FileIOError("File IO Error")
      src_path = os.path.join(work_dir, compile_config['src_name'])
      with open(src_path, "w", encoding="utf-8") as f:
        f.write(src)
      os.chmod(src_path, 0o400)

      Compiler().compileSrc(src_path, work_dir, compile_config)
      exe_path = os.path.join(work_dir, compile_config['exe_name'])
      try:
        os.chmod(exe_path, 0o500)
      except Exception:
        pass
      
      run_config['args'][0] = exe_path
      run_config['time_limit'] = time_limit
      run_config['memory_limit'] = memory_limit

      judge_client = JudgeClient(run_config,
                                work_dir,
                                io_mode,
                                test_case_path,
                                in_file,
                                out_file,
                                is_spj)
      ret = judge_client.run()
      return ret


@app.route('/', defaults={'path': ''})
@app.route('/<path:path>', methods=['POST'])
def server(path):
  if path in ("judge", 'ping'):
    _token = request.headers.get("X-Judge-Server-Token")
    try:
      if _token != token:
        raise TokenVerificationFailed("Invalid token.")
      try:
        data = json.loads(request.data.decode())
      except Exception:
        data = {}
      ret = {"err": None, "data": getattr(JudgeServer, path)(**data)}
    except (CompileError) as e:
      logger.exception(e)
      ret = {"err": e.__class__.__name__, "data": e.message}
    except Exception as e:
      logger.exception(e)
      ret = {"err": "JudgeClientError", "data": e.__class__.__name__ + " :" + str(e)}
  else:
    ret = {"err": "InvalidRequest", "data": "404"}
  return Response(json.dumps(ret), mimetype='application/json')


if DEBUG:
    logger.info("DEBUG=ON")

# gunicorn -w 4 -b 0.0.0.0:8080 server:app
if __name__ == "__main__":
    app.run(debug=DEBUG)
